# today


#!/usr/bin/python3
# coding=utf8
import sys
sys.path.append('/home/pi/MasterPi/')  # 添加路径

import cv2
import time
import math
import signal
import threading
import numpy as np
import yaml_handle
from collections import deque  # 使用双端队列保存轨迹数据
from ArmIK.Transform import *
from ArmIK.ArmMoveIK import *
import HiwonderSDK.Misc as Misc
import HiwonderSDK.Board as Board
from HiwonderSDK.PID import PID  # 导入PID控制器

# ========== 全局配置与初始化 ==========
AK = ArmIK()  # 机械臂逆运动学解算器

# PID控制器初始化（用于循迹偏差控制）
pitch_pid = PID(P=0.28, I=0.16, D=0.25)

# 颜色对应的BGR值（用于可视化绘制）
range_rgb = {
    'red': (0, 0, 255),
    'blue': (255, 0, 0),
    'green': (0, 255, 0),
    'black': (0, 0, 0),
    'white': (255, 255, 255),
}

# 巡线ROI配置（上、中、下三个区域）
# 每个ROI格式：(y_min, y_max, x_min, x_max, weight)，weight为加权系数（下方区域权重更高）
roi_config = [  
    (240, 280,  0, 640, 0.1),  # 上方ROI（低权重）
    (340, 380,  0, 640, 0.3),  # 中间ROI（中权重）
    (430, 460,  0, 640, 0.6)   # 下方ROI（高权重，优先检测）
]
# 计算各ROI的垂直位置偏移（用于坐标映射）
roi_h_list = [
    roi_config[0][0],  # 上方ROI起始行
    roi_config[1][0] - roi_config[0][0],  # 中间ROI高度
    roi_config[2][0] - roi_config[1][0]   # 下方ROI高度
]

lab_data = None  # 颜色阈值数据（从YAML文件加载）
__isRunning = False  # 系统运行标志
__target_color = ('red',)  # 目标检测颜色
template_moments = None     # T型模板的Hu矩

line_centerx = -1  # 线条中心X坐标（-1表示未检测到）
img_centerx = 320  # 图像中心X坐标（640像素宽度时中心为320）
img_size = (640, 480)  # 图像尺寸（宽度, 高度）

# 路口检测参数
CORNER_DETECT_THRESH = 0.5  # 面积突变阈值（当前ROI面积/上一ROI面积 > 阈值时触发检测）
CORNER_TIMEOUT = 3.0         # 转弯超时时间（防止卡死）
CURVATURE_THRESH = 200.0       # 弯道曲率阈值（用于检测路径弯曲程度）
T_JUNCTION_WIDTH_RATIO = 3.0 # T型路口宽度比（下方ROI宽度/中间ROI宽度 > 该值）
T_JUNCTION_OFFSET_THRESH = 80# T型路口横向偏移阈值（中心点偏移超过该值）

corner_detected = False      # 路口检测标志
corner_start_time = 0        # 转弯开始时间戳
turn_direction = 'left'      # 默认转弯方向（可改为'right'）
current_corner_type = None   # 当前路口类型（'t_junction'/'curve'/'straight'）

L_JUNCTION_MIN_FRAMES = 5       # 检测L型所需最小帧数
L_JUNCTION_OFFSET_THRESH = 120  # L型路口横向偏移阈值（像素）
L_JUNCTION_CURVATURE_THRESH = 350  # L型曲率阈值（高于普通弯道）

# 轨迹记录：保存最近5帧的线条中心点，用于曲率计算和弯道检测
center_history = deque(maxlen=5)  

# ========== 工具函数 ==========
def load_config():
    """加载颜色阈值配置（从YAML文件读取LAB颜色范围）"""
    global lab_data
    lab_data = yaml_handle.get_yaml_data(yaml_handle.lab_file_path)

def load_template():
    """加载T型路口模板（用于Hu矩特征匹配）"""
    global template_moments
    template_path = 't_junction_template_red.png'  # T型路口模板图像路径
    template_img = cv2.imread(template_path)
    if template_img is None:
        print(f"警告：未找到模板文件 {template_path}")
        return
    
    # 调整模板尺寸与处理图像一致
    template_img = cv2.resize(template_img, img_size)
    template_lab = cv2.cvtColor(template_img, cv2.COLOR_BGR2LAB)
    
    # 生成红色掩膜
    lower_red = np.array(lab_data['red']['min'], dtype=np.uint8)
    upper_red = np.array(lab_data['red']['max'], dtype=np.uint8)
    template_mask = cv2.inRange(template_lab, lower_red, upper_red)
    
    # 提取最大轮廓并计算Hu矩
    cnts = cv2.findContours(template_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[-2]
    if len(cnts) > 0:
        template_cnt = max(cnts, key=cv2.contourArea)
        template_moments = cv2.HuMoments(cv2.moments(template_cnt)).flatten()
        print("成功加载红色T型模板")
    else:
        print("警告：模板中未检测到红色轮廓")
    

def getAreaMaxContour(contours):
    """筛选最大有效轮廓（过滤小面积噪声）"""
    max_area = 0
    max_contour = None
    for c in contours:
        area = math.fabs(cv2.contourArea(c))  # 计算轮廓面积
        # 过滤面积小于5的噪声，并保留最大轮廓
        if area > max_area and area >= 5:  
            max_area = area
            max_contour = c
    return max_contour, max_area

def MotorStop():
    """停止所有电机（紧急制动）"""
    Board.setMotor(1, 0) 
    Board.setMotor(2, 0)
    Board.setMotor(3, 0)
    Board.setMotor(4, 0)

def setBuzzer(timer):
    """控制蜂鸣器响铃（用于提示检测到路口）"""
    Board.setBuzzer(1)
    time.sleep(timer)
    Board.setBuzzer(0)

# ========== 检测函数 ==========
def detect_t_junction_shape(cnt):
    """基于Hu矩的T型形状匹配"""
    if template_moments is None:
        return False
    
    hu = cv2.HuMoments(cv2.moments(cnt)).flatten()
    # 使用CONTOURS_MATCH_I1距离（值越小越相似）
    similarity = cv2.matchShapes(template_moments, hu, cv2.CONTOURS_MATCH_I1, 0.0)
    return similarity < 0.5  # 相似度阈值（可根据模板调整）

def detect_t_junction(roi_mid, roi_bottom):
    """检测T型路口：通过中间和下方ROI的宽度比和中心点偏移量"""
    if not roi_mid or not roi_bottom:  # 确保两个ROI都有有效数据
        return False
    
    mid_center = roi_mid['center']     # 中间ROI中心点坐标
    bottom_center = roi_bottom['center'] # 下方ROI中心点坐标
    bottom_width = roi_bottom['width']   # 下方ROI轮廓宽度
    mid_width = roi_mid['width']         # 中间ROI轮廓宽度
    
    # 计算宽度比和横向偏移量
    width_ratio = bottom_width / max(mid_width, 1)  # 防止除数为0
    offset = abs(bottom_center[0] - mid_center[0])  # 横向偏移距离
    
    # T型路口判定条件：宽度比大于阈值且偏移量大于阈值
    return width_ratio > T_JUNCTION_WIDTH_RATIO and offset > T_JUNCTION_OFFSET_THRESH

def detect_l_junction(history):
    """检测L型路口：连续偏移+高曲率"""
    if len(history) < L_JUNCTION_MIN_FRAMES:
        return False, 'straight'
    
    # 计算连续帧的横向偏移量
    offsets = [x - img_centerx for x, y in history]
    # 左偏/右偏的连续帧数（需≥3帧）
    left_consec = get_consecutive_count(offsets, lambda o: o < -L_JUNCTION_OFFSET_THRESH)
    right_consec = get_consecutive_count(offsets, lambda o: o > L_JUNCTION_OFFSET_THRESH)
    
    # 计算曲率（复用弯道检测逻辑，但提高阈值）
    curvature = detect_curve(history) and (abs(detect_curve(history)) > L_JUNCTION_CURVATURE_THRESH)
    
    if (left_consec >= 3 or right_consec >= 3) and curvature:
        turn_dir = 'left' if left_consec >= 3 else 'right'
        return True, turn_dir
    else:
        return False, 'straight'

def get_consecutive_count(arr, condition):
    """计算连续满足条件的元素个数"""
    count = 0
    max_count = 0
    for num in arr:
        if condition(num):
            count += 1
            max_count = max(max_count, count)
        else:
            count = 0
    return max_count

def detect_curve(history):
    """检测弯道：通过中心点轨迹的曲率变化（二阶差分）"""
    if len(history) < 5:  # 至少需要5帧数据计算曲率
        return False
    
    # 计算相邻帧的横向偏移量
    dx = [history[i][0] - history[i-1][0] for i in range(1, len(history))]
    # 计算曲率（偏移量的变化率，近似二阶导数）
    curvature = abs(dx[2] - dx[1])
    print(f'Curvature: {curvature}')
    return curvature > CURVATURE_THRESH  # 曲率超过阈值则判定为弯道

# ========== 硬件控制函数 ==========
def initMove():
    """机械臂与舵机初始位置（开机复位）"""
    Board.setPWMServoPulse(1, 1500, 800)  # 设置舵机初始位置
    # 机械臂初始姿态（肩、肘、腕关节角度）
    AK.setPitchRangeMoving((0, 7, 11), -60, -90, 0, 1500)  
    MotorStop()  # 停止所有电机

def reset():
    """重置全局状态（用于系统启动或重启）"""
    global line_centerx, __target_color, corner_detected, current_corner_type, center_history
    line_centerx = -1          # 清除线条中心坐标
    __target_color = ()         # 清除目标颜色
    corner_detected = False     # 清除路口检测标志
    current_corner_type = 'straight'  # 初始状态为直行
    center_history.clear()      # 清空轨迹记录

def init():
    """系统初始化（加载配置+硬件复位）"""
    print("VisualPatrol Init")
    load_config()  # 加载颜色阈值配置
    initMove()     # 硬件复位

def start():
    """启动系统（设置运行标志+重置状态）"""
    global __isRunning
    reset()
    __isRunning = True
    print("VisualPatrol Start")

def stop():
    """停止系统（关闭运行标志+制动电机）"""
    global __isRunning
    __isRunning = False
    MotorStop()
    print("VisualPatrol Stop")

def exit():
    """退出系统（关闭运行标志+制动电机）"""
    global __isRunning
    __isRunning = False
    MotorStop()
    print("VisualPatrol Exit")

# ========== 电机控制线程（核心控制逻辑） ==========
def move():
    """电机控制线程：负责循迹和路口转向"""
    global line_centerx, corner_detected, corner_start_time, current_corner_type, current_turn_direction
    
    while True:
        if __isRunning:
            if line_centerx != -1:  # 检测到有效线条时执行循迹
                # 计算循迹偏差（线条中心与图像中心的横向距离）
                error = line_centerx - img_centerx  

                # 增加死区处理：微小误差不响应，减少震荡
                if abs(error) < 30:
                    error = 0

                # PID控制器更新输出
                pitch_pid.update(error)
                pid_output = pitch_pid.output
                # 限幅处理（防止电机速度超过±100）
                pid_output = max(-100, min(pid_output, 100))  
                # 将PID输出映射到电机速度范围（-50到50）
                base_speed = Misc.map(pid_output, -100, 100, -50, 50)
                
                if corner_detected:  # 检测到路口时执行转向逻辑  
                    # 根据路口类型执行不同转向策略
                    if current_corner_type == 't_junction':
                        # T型路口转向逻辑
                        if time.time() - corner_start_time < 3.0:  # T型路口转向时间
                            Board.setMotor(1, -40)  
                            Board.setMotor(2, 40)
                            Board.setMotor(3, -40)
                            Board.setMotor(4, 40)
                        else:
                            corner_detected = False
                            current_corner_type = 'straight'
                            setBuzzer(0.5)  # 提示转弯完成
                            pitch_pid.I = 0
                            pitch_pid.D = 0
                            pitch_pid.last_error = 0

                    elif current_corner_type == 'l_junction':
                        # L型路口转向逻辑
                        if time.time() - corner_start_time < 2.5:  # L型路口转向时间
                            if current_turn_direction == 'left':
                                Board.setMotor(1, -35)  
                                Board.setMotor(2, 35)
                                Board.setMotor(3, -35)
                                Board.setMotor(4, 35)
                            else:
                                Board.setMotor(1, 35)  
                                Board.setMotor(2, -35)
                                Board.setMotor(3, 35)
                                Board.setMotor(4, -35)
                        
                        else:
                            corner_detected = False
                            current_corner_type = 'straight'
                            setBuzzer(0.5)
                            pitch_pid.I = 0
                            pitch_pid.D = 0
                            pitch_pid.last_error = 0

                    elif current_corner_type == 'curve':
                        # 弯道控制：降低基础速度，增大转向灵敏度
                        base_speed = Misc.map(pid_output, -100, 100, -30, 30)
                        Board.setMotor(1, int(20 - base_speed))
                        Board.setMotor(2, int(20 + base_speed))
                
                
                else:  # 正常循迹时的电机控制
                    Board.setMotor(1, int(40 - base_speed))  # 左电机速度
                    Board.setMotor(2, int(40 + base_speed))  # 右电机速度
                    Board.setMotor(3, int(40 - base_speed))
                    Board.setMotor(4, int(40 + base_speed))
            else:  # 未检测到线条时停止电机
                MotorStop()
                time.sleep(0.01)
        else:
            time.sleep(0.01)

# 启动电机控制线程（守护线程，主程序退出时自动终止）
th_move = threading.Thread(target=move)
th_move.setDaemon(True)
th_move.start()

# ========== 图像处理核心逻辑（视觉检测） ==========
def run(img):
    """图像处理主函数：负责ROI检测、轮廓分析和路口识别"""
    global line_centerx, corner_detected, corner_start_time, current_corner_type, center_history, template_moments
    
    red_t_detected = False  # 红色T型路口检测标志
    
    if not __isRunning or not __target_color:  # 系统未运行或未设置目标颜色时直接返回
        return img

    img_copy = img.copy()
    img_h, img_w = img.shape[:2]
    # 调整图像尺寸到640x480（与ROI配置匹配）
    frame = cv2.resize(img_copy, img_size, interpolation=cv2.INTER_NEAREST)
    # 高斯模糊去噪
    frame = cv2.GaussianBlur(frame, (3, 3), 3)  
    # 颜色空间转换：BGR -> LAB（对光照变化更鲁棒）
    frame_lab = cv2.cvtColor(frame, cv2.COLOR_BGR2LAB)

    centroid_x_sum = 0  # 加权中心点X坐标累加
    weight_sum = 0      # 权重累加
    prev_roi_area = 0   # 上一ROI的最大轮廓面积（用于面积突变检测）
    roi_mid_info = {}   # 中间ROI的信息（中心点、宽度）
    roi_bottom_info = {}# 下方ROI的信息（中心点、宽度）

    for i, (y_min, y_max, x_min, x_max, weight) in enumerate(roi_config):
        # 提取当前ROI区域
        roi = frame_lab[y_min:y_max, x_min:x_max]
        mask = None  # 颜色掩膜（用于筛选目标颜色）
        
        # 遍历所有目标颜色，生成复合掩膜
        for color in __target_color:
            if color not in lab_data:
                continue
            # 获取LAB颜色阈值
            lower = np.array(lab_data[color]['min'], dtype=np.uint8)
            upper = np.array(lab_data[color]['max'], dtype=np.uint8)
            # 生成当前颜色的掩膜
            _mask = cv2.inRange(roi, lower, upper)
            # 形态学闭运算：填充小孔、连接邻近区域
            kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
            _mask = cv2.morphologyEx(_mask, cv2.MORPH_CLOSE, kernel)
            
            if mask is None:
                mask = _mask
            else:
                mask = cv2.bitwise_or(mask, _mask)  # 合并多颜色掩膜
            

        # 如果当前ROI没有有效掩膜，则跳过后续处理
        if mask is None:  # 无有效掩膜时跳过
            continue
        
        # 提取轮廓并筛选最大有效轮廓
        cnts = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[-2]
        cnt, area = getAreaMaxContour(cnts)
        
        if cnt is not None:
            # 获取轮廓的最小外接矩形
            rect = cv2.minAreaRect(cnt)
            box = np.int0(cv2.boxPoints(rect))
            (cx, cy), (cw, ch), angle = rect  # 中心点、宽高、旋转角度
            
            # 将ROI内的坐标映射回原始图像尺寸
            box[:, 1] += sum(roi_h_list[:i]) + roi_config[0][0]
            box[:, 1] = Misc.map(box[:, 1], 0, img_size[1], 0, img_h)
            box[:, 0] = Misc.map(box[:, 0], 0, img_size[0], 0, img_w).astype(int)
            
            # 绘制轮廓（红色矩形）
            cv2.drawContours(img, [box], -1, range_rgb['red'], 2)
            # 计算中心点并绘制（绿色圆点）
            center_x = cx + x_min
            center_y = cy + y_min
            center_x = Misc.map(center_x, 0, img_size[0], 0, img_w)
            center_y = Misc.map(center_y, 0, img_size[1], 0, img_h)
            cv2.circle(img, (int(center_x), int(center_y)), 5, range_rgb['green'], -1)
            
            # 加权求和计算最终中心点
            centroid_x_sum += center_x * weight
            weight_sum += weight
            prev_roi_area = area  # 更新上一ROI面积
            
            # 记录中间和下方ROI的信息（用于T型路口检测）
            if i == 1:  # 中间ROI（索引1）
                roi_mid_info = {
                    'center': (center_x, center_y),
                    'width': cw
                }
            elif i == 2:  # 下方ROI（索引2）
                roi_bottom_info = {
                    'center': (center_x, center_y),
                    'width': cw
                }
            
            center_x_filtered = np.median([c[0] for c in center_history])
            
            # 保存中心点历史（用于弯道检测）
            center_history.append((center_x, center_y))
    

            # 如果当 # 红色T型形状匹配（仅在下方ROI检测）
            if i == 2 and color == 'red':
                if detect_t_junction_shape(cnt):
                    red_t_detected = True
                    cv2.putText(img, "T-Junction Detected", (int(center_x)-50, int(center_y)), 
                                cv2.FONT_HERSHEY_SIMPLEX, 0.6, range_rgb['green'], 2)
    # ------ 路口类型检测与逻辑处理 ------
    if len(center_history) > 0:
        # 计算最终线条中心点（加权平均）
        line_centerx = int(centroid_x_sum / weight_sum) if weight_sum else -1
        
        if abs(line_centerx - img_centerx) < 100:
            current_corner_type = 'straight'
            current_turn_direction = 'straight'
            corner_detected = False
        
        # 优先检测T型路口
        if red_t_detected or detect_t_junction(roi_mid_info, roi_bottom_info):
            current_corner_type = 't_junction'
            current_turn_direction = 'left'
            corner_detected = True
            corner_start_time = time.time()
            setBuzzer(0.3)
            print("检测到T型路口！")
        
        # 其次检测L型路口
        else:
            l_detected, turn_dir = detect_l_junction(center_history)
            if l_detected:
                current_corner_type = 'l_junction'
                current_turn_direction = turn_dir
                corner_detected = True
                corner_start_time = time.time()
                setBuzzer(0.2)
                print(f"检测到L型路口，方向：{turn_dir}")
            else:
                if len(center_history) >= 5 and detect_curve(list(center_history)) and not corner_detected:
                    current_corner_type = 'curve'
                    corner_detected = False  # 弯道不强制转向，仅循迹调整
                    print("检测到弯道！")
                else:
                    current_corner_type = 'straight'
                    corner_detected = False  # 确保标志清除
    
        # 在图像上绘制检测到的路口类型（用于调试可视化）
    if current_corner_type != 'straight':
        cv2.putText(img, f"Detected: {current_corner_type}", (20, 40), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, range_rgb['blue'], 2)
    
    return img

# ========== 主程序入口 ==========
def Stop():
    """信号处理函数：捕获Ctrl+C并安全停止系统"""
    global __isRunning
    __isRunning = False
    print('关闭中...')
    MotorStop()  # 紧急停止电机

# 用于外部调用的函数

def patrol_run(obstacle):
    """
    外部调用函数，
    """
    global __target_color
    
    init()
    start()
    cap = cv2.VideoCapture('http://127.0.0.1:8080?action=stream')
    __target_color = ('red',)
    while True:
        if not obstacle:
            __isRunning = True
            while __isRunning:
                    ret,img = cap.read()
                    if ret:
                        # 处理图像（检测线条和路口）
                        output_frame = run(img)
                        # 调整显示尺寸（便于观察）
                        display = cv2.resize(output_frame, (320, 240))
                        # cv2.imshow('VisualPatrol', display)  # 延迟高可尝试注释此行
                        # 按ESC键退出
                        if cv2.waitKey(1) == 27:
                            break
                    else:
                        time.sleep(0.01)
            stop()
        else:
            continue


if __name__ == '__main__':
    init()  # 系统初始化
    start()  # 启动系统
    
    # 注册信号处理（捕获Ctrl+C）
    signal.signal(signal.SIGINT, Stop)
    
    # 初始化摄像头（根据实际流地址调整）
    cap = cv2.VideoCapture('http://127.0.0.1:8080?action=stream')
    __target_color = ('red',)  # 设置目标检测颜色（如黑色线条）
    
    # 主循环：持续捕获图像并处理
    while __isRunning:
        ret, img = cap.read()
        if ret:
            # 处理图像（检测线条和路口）
            output_frame = run(img)
            # 调整显示尺寸（便于观察）
            display = cv2.resize(output_frame, (320, 240))
            cv2.imshow('VisualPatrol', display)
            # 按ESC键退出
            if cv2.waitKey(1) == 27:
                break
        else:
            time.sleep(0.01)  # 读取失败时短暂休眠
    
    # 释放资源
    cv2.destroyAllWindows()
    cap.release()
    stop()  # 停止系统